from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from typing import Type


# print(int, float)
class MultipleTool(BaseTool):
    name: str = "乘法工具"
    description: str = "用于两个数相乘"

    class InputSchema(BaseModel):
        number1: float = Field(description="两数相乘，第一个参数")
        number2: float = Field(description="两数相乘，第二个参数")

    args_schema: Type[BaseModel] = InputSchema

    # def _run(self, number1: float, number2: float):
    #     return number1 * number2

    def _run(self, number1: int, number2: int):
        return number1 * number2


multiple_tool = MultipleTool()

# result = multiple_tool.invoke(input={"number1": 2, "number2": 3})
result = multiple_tool.invoke({"number1": 2, "number2": 3})
print(result)
print(multiple_tool.name)
print(multiple_tool.description)
